import numpy as np

"""BBE32EP CMDS
"""
def BBE_MUL16PACKQ(a, b): return ((a.astype('i') * b.astype('i')) >> 15).astype('h')
def BBE_MUL16PACKL(a, b): return ((a.astype('i') * b.astype('i')) & 0xFFFF).astype('H')
def BBE_ADDSR1RNX16(a, b): return np.round((a.astype('i') + b.astype('i'))/2).astype('h')
def BBE_SVNX16F_I(src, dst, b): return np.where(b, dst, src)

def cmplx2int16(inArray: np.ndarray):
    array = inArray.ravel().astype(np.complex128)
    array.dtype = np.float64
    array = array / np.max(abs(array)) * 0x7FFF
    array = np.round(array, 0).astype('h')
    return array

def SptHist(mag: np.ndarray, inputType='real'):
    bins = 256*np.arange(65)  # log2 coder
    ref = bins[1:] - 128
    if inputType == 'real':
        mag = np.round(20*np.log10(mag)*85).astype('H')
        ref = 10**(ref / 85 / 20)
    elif inputType == 'dB':
        mag = (mag * 85).astype('H')
        ref = ref / 85
    else:
        mag = mag.astype('H')
    ret = []
    for i in mag:
        tmp = np.histogram(i, bins)
        ret.append(ref[np.argmax(tmp[0])])
    ret = np.array(ret)
    return ret

def pyHist(mag: np.ndarray, inputType='real'):
    bins = 256*np.arange(65)  # log2 coder
    ref = bins[1:] - 128
    if inputType == 'real':
        mag = np.round(20*np.log10(mag)*85).astype('H')
        ref = 10**(ref / 85 / 20)
    elif inputType == 'dB':
        mag = (mag * 85).astype('H')
        ref = ref / 85
    else:
        mag = mag.astype('H')
    ret = []
    for i in mag:
        tmp = np.histogram(i, bins)
        idx = np.argmax(tmp[0])
        if ((idx > 0) & (idx < len(tmp[0]) - 1)):
            diff = (tmp[0][idx-1] + tmp[0][idx+1] - 2 * tmp[0][idx]) / 2
            fIdx = 0 if diff==0 else (tmp[0][idx-1] - tmp[0][idx+1]) / diff
            idx = idx + fIdx
        ret.append((idx + 0.5) * 256)
    ret = np.array(ret)
    return ret

def SptSva(xkReal, xkImag):
    """ this is only a fixed point flow test for simulate bbe32 dsp
    """
    zero = 0
    xkReal = xkReal.astype('h')
    xkImag = xkImag.astype('h')
    l = np.roll(np.arange(len(xkReal)), -1)
    r = np.roll(np.arange(len(xkReal)), 1)
    mReal = BBE_ADDSR1RNX16(xkReal[l], xkReal[r])
    mImag = BBE_ADDSR1RNX16(xkImag[l], xkImag[r])
    a_n = xkReal.astype('q') * mReal.astype('q') + xkImag.astype('q') * mImag.astype('q')
    a_d = mReal.astype('q') * mReal.astype('q') + mImag.astype('q')* mImag.astype('q')
    stage0 = a_n > zero
    stage2 = a_n < a_d
    stage1 = ~(stage0 & stage2)
    a_rem = np.where(stage1, 0, a_n / a_d * 0x7fff).astype('h')
    xkReal = BBE_SVNX16F_I(xkReal - mReal, xkReal, stage2)
    xkImag = BBE_SVNX16F_I(xkImag - mImag, xkImag, stage2)
    xkReal = BBE_SVNX16F_I(xkReal - BBE_MUL16PACKQ(mReal, a_rem), xkReal, stage1)
    xkImag = BBE_SVNX16F_I(xkImag - BBE_MUL16PACKQ(mImag, a_rem), xkImag, stage1)
    return xkReal, xkImag

def SptSva2(xkReal, xkImag):
    """ this is only a fixed point flow test for simulate bbe32 dsp
    """
    xkReal = xkReal.astype('f')
    xkImag = xkImag.astype('f')
    l = np.roll(np.arange(len(xkReal)), -1)
    r = np.roll(np.arange(len(xkReal)), 1)
    mReal = np.round((xkReal[l] + xkReal[r]) / 2)
    mImag = np.round((xkImag[l] + xkImag[r]) / 2)
    a_n = xkReal * mReal + xkImag * mImag
    a_d = mReal * mReal + mImag * mImag
    stage0 = a_n < 0
    stage2 = a_n > a_d
    stage1 = ~(stage0 | stage2)
    a = a_n / a_d
    xkReal = np.where(stage1, xkReal - a * mReal, xkReal)
    xkImag = np.where(stage1, xkImag - a * mImag, xkImag)
    xkReal = np.where(stage2, xkReal - mReal, xkReal)
    xkImag = np.where(stage2, xkImag - mImag, xkImag)
    return xkReal, xkImag

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    import time
    from utilsCommon import pySva
    #%% FFT test
    RBIN = 4096
    sig = 20*np.exp(50.685j*np.arange(RBIN)/RBIN*2*np.pi) + np.random.random(RBIN) - 0.5
    raw = cmplx2int16(np.fft.fft(sig))>>1
    svaRe, svaIm = SptSva(raw[::2], raw[1::2])
    sva = svaRe + 1j * svaIm
    plt.figure()
    plt.subplot(211)
    plt.title('FFT Test')
    plt.plot(20*np.log10(abs(raw[::2] + 1j*raw[1::2])), label='rect win')
    # plt.plot(20*np.log10(abs(np.dot(steerVec, sig))), label='rect win float')
    plt.plot(20*np.log10(abs(sva)), label='SVA')
    plt.legend()
    plt.subplot(212)
    plt.plot(20*np.log10(abs(np.fft.fft(sig))), label='rect win')
    # plt.plot(20*np.log10(abs(np.dot(steerVec, sig))), label='rect win float')
    plt.plot(20*np.log10(abs(pySva(np.fft.fft(sig)))), label='SVA')
    plt.show()
    #%% DFT test
    snap = 8
    N = 194
    nfft = 512
    p = np.r_[0,np.arange(2,192),193]
    snapSig = []
    for i in range(snap):
        sig = 10*np.exp(50.25j*(np.arange(N)/N*2*np.pi + np.random.random(1) * 2 * np.pi)) + (np.random.random(N) - 0.5)
        sig = sig[p]
        snapSig.append(sig)
    az = np.arange(nfft)-nfft//2
    steerVec = np.outer(az/nfft*2*np.pi, p)
    steerVec = np.exp(-1j*steerVec)
    raw = cmplx2int16(np.dot(steerVec, sig))
    svaRe, svaIm = SptSva(raw[::2], raw[1::2])
    sva = svaRe + 1j * svaIm
    plt.figure()
    plt.title('DFT Test')
    plt.plot(20*np.log10(abs(raw[::2]+ 1j*raw[1::2])), label='rect win')
    # plt.plot(20*np.log10(abs(np.dot(steerVec, sig))), label='rect win float')
    plt.plot(20*np.log10(abs(sva)), label='SVA')
    plt.legend()
    plt.show()

